import jpcap.*;
import jpcap.packet.Packet;
import jpcap.packet.ARPPacket;
import jpcap.packet.IPPacket;

//import java.io.FileWriter;
import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;

/**
 * Project: SecuritySnort
 * This class takes care of capturing packets and either saving
 * them or displaying them in the console, or doing both.
 *
 * @author Jonathan Buttner
 * @version Jan 23, 2010
 */
public class CaptureCmd extends Command
{
	/**************************************************************************
	 * 							Globals
	 *************************************************************************/
	private NetworkInterface _device;
	private JpcapCaptor _captor;
	private PacketThread _pThread;
	private OutputTo _save;
	private JpcapWriter _writer;
	private String _fileName;
	//private FileWriter _fileWriter;
	private IPSaver _tallyIP;
	
	/**************************************************************************
	 * 							Constructors
	 *************************************************************************/
	/**
	 * This constructor sets up the required objects to allow listening to
	 * packets from the device.
	 * 
	 * @param saver is the IP collector that keeps track of what
	 * IP's have been seen while capturing.
	 */
	public CaptureCmd( IPSaver saver )
	{
		//_fileWriter = null;
		_fileName = "";
		_writer = null;
		_save = OutputTo.DISP;
		_device = null;
		_captor = null;
		_pThread = null;
		_name = "Listen";
		//_tallyIP = new IPSaver();
		_tallyIP = saver;
	}
	
	/**************************************************************************
	 * 							Public Methods
	 *************************************************************************/
	/**
	 * This method will capture packets for a given device.
	 * 
	 * @param command is the tokenized version of the command given
	 * by the user.
	 */
	public void exeCommand( ArrayList < String > cmdStrs )
	{
		if ( _device == null )
		{
			System.out.println("Error! The device was not set correctly!"
					+ "\nPlease report this message to the"
					+ "programs so they can fix it!");
			return;
		}
		try
		{
			if ( cmdStrs.size() == 1 )
			{
				_captor = JpcapCaptor.openDevice( _device, 65535, true, 20 );
				_pThread = new PacketThread( -1 );
			}
			else
			{
				try
				{
					int numPackets = Integer.parseInt( cmdStrs.get(1) );
					_captor = JpcapCaptor.openDevice( _device, 65535, true, 20 );
					_pThread = new PacketThread( numPackets );
				}
				catch( NumberFormatException e )
				{
					System.out.println( "The number of packets to capture"
							+ "was an invalid number.\nCommand: 'listen' failed!");
					return;
				}
			}			
		}
		catch( IOException e )
		{
			System.out.println( "SecuritySnort was unable to open"
					+ " the selected device.\nCommand: 'listen' failed!" );
			return;
		}
		
		try
		{
			if ( _save == OutputTo.BOTH || _save == OutputTo.FILE )
			{
				_writer = JpcapWriter.openDumpFile(_captor, _fileName );
				//_fileWriter = new FileWriter( _fileName + "1" );
			}
		}
		catch( IOException e )
		{
			System.out.println( "SecuritySnort was unable to write"
					+ " to the specified file:\n" + _fileName 
					+ "\nCommand: 'listen' failed!" );
			return;
		}
		
		Thread newThread = new Thread( _pThread );
		newThread.start();
	}
	
	/**
	 * This method stops the execution of the listening thread and command.
	 */
	public void exeStop()
	{
		if ( _pThread != null )
			_pThread.stop();
	}
	
	/**************************************************************************
	 * 							Getter/Setter Methods
	 *************************************************************************/
	/**
	 * This method returns the NetworkInterface object being listened to.
	 * 
	 * @return a NetworkInterface object being used to capture packets.
	 */
	public NetworkInterface getDevice()
	{
		return _device;
	}
	
	/**
	 * This method sets the NetworkInterface object (device) to capture
	 * packets from. NOTE: Set this before the listen is begun.
	 * 
	 * @param device is a NetworkInterface object used to capture packets.
	 */
	public void setDevice( NetworkInterface device )
	{
		_device = device;
	}
	
	/**
	 * This method allows the user to specify
	 * where the packets should be displayed.
	 * 
	 * @param output is the switch representing
	 * where the packets are displayed.
	 */
	public void setOutput( OutputTo output )
	{
		_save = output;
	}
	
	/**
	 * This method sets the file name that the packets
	 * will be saved to after they are captured.
	 * 
	 * @param name is the string representing the file name.
	 */
	public void setFilename( String name )
	{
		_fileName = name;
	}
	
	/**************************************************************************
	 * 							Private Methods
	 *************************************************************************/
	/**
	 * This method will try to construct a InetAddress object to make
	 * storage in the tree more consistent.
	 * 
	 * @param ip is the IP Address that should be stored in the InetAddress
	 * @return an InetAddress object or null if it failed.
	 */
	private InetAddress getInetAdd( String ip )
	{
		try
		{
			return InetAddress.getByName( ip );
		} 
		catch (UnknownHostException e)
		{
			return null;
		}
	}
	
	/**
	 * This method determines if a given make a mac address that is unknown.
	 * Aka it tests to see if it is all 0s or Fs.
	 * 
	 * @param mac is a String to test.
	 * @return true if the mac address is unknown.
	 */
	private boolean isUnknownMAC( String mac )
	{
		if ( mac.equals( "00:00:00:00:00:00" )
				|| mac.equalsIgnoreCase( "FF:FF:FF:FF:FF:FF" ) )
			return true;
		else
			return false;
	}
	
	/**
	 * This method takes care of saving the IP and MAC Address that is
	 * contained in the receieved packet.
	 * 
	 * @param packy is the packet to extract the information from.
	 */
	private void saveIPMAC( Packet packy )
	{
		if ( packy instanceof ARPPacket )
		{
			ARPPacket arp = (ARPPacket)packy;
							
			if ( !isUnknownMAC(arp.getSenderHardwareAddress().toString()))
			{
				String senderIP = 
					arp.getSenderProtocolAddress().toString().replaceAll("/", "");
				InetAddress inetSend = getInetAdd( senderIP );
				if ( inetSend != null )
					_tallyIP.addIP(inetSend, arp.getSenderHardwareAddress() );
				else
					_tallyIP.addIP(senderIP, arp.getSenderHardwareAddress().toString());
			}
			
			if ( !isUnknownMAC(arp.getTargetHardwareAddress().toString()))
			{
				String recIP = 
					arp.getTargetProtocolAddress().toString().replaceAll("/", "");
				InetAddress inetRec = getInetAdd( recIP );
				//_tallyIP.addIP(inetRec, arp.getTargetHardwareAddress() );
				
				if ( inetRec != null )
					_tallyIP.addIP(inetRec, arp.getTargetHardwareAddress() );
				else
					_tallyIP.addIP(recIP, arp.getTargetHardwareAddress().toString());
			}
		}
		else if ( packy instanceof IPPacket )
		{
			_tallyIP.addIP( ((IPPacket)packy).dst_ip, "" );
			_tallyIP.addIP( ((IPPacket)packy).src_ip, "" );
		}
	}
	/**************************************************************************
	 * 							Private Classes
	 *************************************************************************/
	/**
	 * This class is used to create a new thread to run the packet capturer on.
	 */
	private class PacketThread implements Runnable
	{
		private int _nPackets;
		
		/**
		 * This constructor sets up the PacketThread object.
		 * 
		 * @param numPackets is the number of packets to capture.
		 */
		public PacketThread( int numPackets )
		{
			_nPackets = numPackets;
		}
		
		/**
		 * This method stops the thread and capturing.
		 */
		public void stop()
		{
			_captor.breakLoop(); //NOTE TO SELF: Never call this on a closed JpcapCaptor, JVM will crash and burn!
			_captor.setPacketReadTimeout( 10 ); //just in case breakLoop doesn't work...
			_captor.close();
			if (_writer != null )
			{
				/*try
				{
					_fileWriter.close();
				}
				catch (IOException e )
				{
					System.out.println( "Error while trying to"
							+ "stop the file writer" );
				}*/
				_writer.close();
			}
			_pThread = null;
		}
		
		/**
		 * This method is run as soon as the thread is created.
		 */
		public void run()
		{
			if ( _save == OutputTo.FILE )
				_captor.loopPacket( _nPackets, new PacketGetterFile() );
			else if ( _save == OutputTo.BOTH )
				_captor.loopPacket( _nPackets, new PacketGetterBoth() );
			else
				_captor.loopPacket( _nPackets, new PacketGetterDisp() );
			
			System.out.println( "\nSecuritySnort has finished"
					+ " caputuring packets." );
		}
	}
	
	/**
	 * This class will display a packet to the screen when it is received.
	 */
	private class PacketGetterDisp implements PacketReceiver
	{
		/**
		 * This method is run everytime a packet is received.
		 */
		public void receivePacket( Packet packy )
		{
			System.out.println( packy );
			saveIPMAC( packy );
		}	
	}
	
	/**
	 * This class will save a packet to a file when it is received.
	 */
	private class PacketGetterFile implements PacketReceiver
	{
		
		/**
		 * This method is run everytime a packet is received.
		 * 
		 * @param packy is the packet that was receieved/
		 */
		public void receivePacket( Packet packy )
		{
			//String data = new String( packy.data );
			/*try
			{
				_fileWriter.write( "Data: " + data );
			} catch (IOException e)
			{
				// TODO Auto-generated catch block
				e.printStackTrace();
			}*/
			/*if ( packy instanceof IPPacket )
			{
				_tallyIP.addIP( ((IPPacket)packy).dst_ip, "" );
				_tallyIP.addIP( ((IPPacket)packy).src_ip, "" );
			}
			else if ( packy instanceof ARPPacket )
			{
				
			}*/
			saveIPMAC( packy );
			_writer.writePacket( packy );
		}
	}
	
	/**
	 * This class will save a packet to a file and display it on the screen.
	 */
	private class PacketGetterBoth implements PacketReceiver
	{
		/**
		 * This method is run everytime a packet is received.
		 * It will save a packet to a file and display it on the screen.
		 * 
		 * @param packy is the packet that was received.
		 */
		public void receivePacket( Packet packy )
		{
			saveIPMAC( packy );
			_writer.writePacket( packy );
			System.out.println( packy );
		}
	}
}
